# -*- coding: utf-8 -*-
"""
@Author: Goware
@Date: 2024/11/20
@Description: 序列化
"""
import os
import re
from itertools import zip_longest

from qiskit.qasm3 import dump

from optimize_dp_table.new_dp_algorithm import NewDpAlgorithmProgram
from optimize_dp_table.setting import DPTABLES_TYPE


class SerializedProgram:
    """
    对DP表进行序列化,向控制台打印量子电路、张量链表示并导出

    Attributes:
        dp_tables (list): 生成的DP表

    Methods:
        construct_circuit(): 根据dp_tables生成指定数量子电路
        construct_tensor_chain(): 根据dp_tables表生成张量链表示
        _apply_gate(): 根据操作类型向量子电路中添加对应的门
        export_qasm(): 序列化 DP 表生成的量子电路并保存为 QASM 文件。
        print_circuit(): 打印生成的量子电路到控制台；生成张量链并表示
    """

    def __init__(self, dp_tables):
        # 初始化 DP 表类型和 DP 表
        self.dp_types = DPTABLES_TYPE
        self.dp_tables = dp_tables

    def preprocess_dp_table(self):
        """
        将 DP 表统一为张量形状列表，支持一维和二维输入格式。
        :param dp_table: DP 表，可能是一维或二维的张量列表。
        :return: 统一的张量形状列表 [(rows, cols), ...]。
        """
        if not self.dp_tables:
            raise ValueError("DP table is empty or invalid.")
        tensor_shapes = []
        # 判断是否为一维 DP 表：每个元素是一个元组 (rows, cols)
        if all(isinstance(tensor, tuple) and len(tensor) == 2 for tensor in self.dp_tables):
            # 直接使用一维 DP 表的张量形状
            tensor_shapes = self.dp_tables
        # 判断是否为二维 DP 表：每个元素是一个列表（表示张量的行）
        elif all(isinstance(tensor, list) and isinstance(tensor[0], list) for tensor in self.dp_tables):
            for tensor in self.dp_tables:
                rows = len(tensor)  # 行数
                cols = max(len(row) for row in tensor) if tensor else 0  # 最大列数
                tensor_shapes.append((rows, cols))
        else:
            raise ValueError("Unsupported DP table format. DP table must be one-dimensional or two-dimensional.")
        return tensor_shapes

    def construct_circuit(self):
        """
        正确构建量子电路：按时间步（列）优先遍历，每一列为一个时间层
        """
        from qiskit import QuantumCircuit

        qubits = sorted({q for tbl in self.dp_tables for q in tbl.qubits})
        qid2row = {q: i for i, q in enumerate(qubits)}
        circuit = QuantumCircuit(len(qubits))

        # 2. 列（时间步）优先遍历
        max_cols = max(len(row) for tbl in self.dp_tables for row in tbl)
        for col in range(max_cols):
            for tbl in self.dp_tables:
                for local_row, global_q in enumerate(sorted(tbl.qubits)):
                    if col >= len(tbl[local_row]):
                        continue
                    op = tbl[local_row][col]
                    if op:  # 空串跳过
                        gate = op.lower()
                        self._apply_gate(circuit, qid2row[global_q], gate,  len(qubits))

        return circuit

    def construct_tensor_chain(self):
        """
        根据 DP 表生成张量链表示。
        """
        tensor_chain = []
        a = NewDpAlgorithmProgram(self.dp_tables)

        layers = a.split_into_layers_in_chain_order()

        interleaved = [tbl
                       for tup in zip_longest(*layers)   # 交错
                       for tbl in tup
                       if tbl is not None]

        for tbl in interleaved:
            for local_row, global_q in enumerate(sorted(tbl.qubits)):
                for op in tbl[local_row]:
                    if op:
                        tensor_chain.append((op.lower(), (global_q,)))

        return tensor_chain

    def _apply_gate(self, circuit, qubit_idx, operation, num_qubits):
        """
        支持带参数的 gate（如 ry(3.14)）以及常规 gate
        """

        if isinstance(operation, list):
            operation = operation[0] if operation else ''
        operation = operation.strip()

        # 参数门如 ry(3.14)
        match = re.match(r"(r[xyz])\(([^)]+)\)", operation)
        if match:
            gate_type, param = match.groups()
            theta = float(param)
            if gate_type == "ry":
                circuit.ry(theta, qubit_idx)
            elif gate_type == "rx":
                circuit.rx(theta, qubit_idx)
            elif gate_type == "rz":
                circuit.rz(theta, qubit_idx)
            return

        # 多 qubit 门
        if operation.startswith("cx") or operation.startswith("cz") or operation.startswith(
                "cy") or operation.startswith("swap"):
            # 暂时假设总是作用于 q[i] 和 q[i+1]
            if qubit_idx < num_qubits - 1:
                if operation == "cx":
                    circuit.cx(qubit_idx, qubit_idx + 1)
                elif operation == "cz":
                    circuit.cz(qubit_idx, qubit_idx + 1)
                elif operation == "cy":
                    circuit.cy(qubit_idx, qubit_idx + 1)
                elif operation == "swap":
                    circuit.swap(qubit_idx, qubit_idx + 1)
            return

        # 单比特门
        gate_map = {
            "h": circuit.h,
            "x": circuit.x,
            "y": circuit.y,
            "z": circuit.z,
            "s": circuit.s,
            "t": circuit.t,
        }
        if operation in gate_map:
            gate_map[operation](qubit_idx)

    def export_qasm(self, filename):
        """
        序列化 DP 表生成的量子电路并保存为 QASM 文件。
        """
        try:
            output_dir = "qasm_files"
            os.makedirs(output_dir, exist_ok=True)
            # 构建量子电路并验证
            circuit = self.construct_circuit()
            # 序列化为 QASM 文件
            filepath = os.path.join(output_dir, filename)  # 拼接文件路径
            with open(filepath, "w") as f:
                # 导入时请使用 QASM的dump 格式化函数
                dump(circuit, f)
            print(f"QASM file generated successfully: {filepath}")
        except Exception as e:
            print(f"Failed to write QASM file: {e}")

    def print_circuit(self):
        """
        打印生成的量子电路到控制台；生成张量链并表示
        """
        # 构建电路
        circuit = self.construct_circuit()
        # 打印电路
        print(circuit.draw())
        # 构建张量链
        tensor_chain = self.construct_tensor_chain()
        # 打印张量链表示
        print("\nTensor Chain Representation:")
        # 遍历张量链
        for tensor in tensor_chain:
            # 获取操作和量子比特
            operation, qubits = tensor
            # 将量子比特转换为字符串
            qubits_str = ", ".join(map(str, qubits))
            # 打印张量
            print(f"Tensor(Operation={operation}, Qubits=({qubits_str}))")

